/*-------------------------------------------------------------------------
 *
 * auto_explain.c
 *
 *
 * Copyright (c) 2008-2012, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *	  contrib/auto_explain/auto_explain.c
 *
 *-------------------------------------------------------------------------
 */
#include "postgres.h"
#include "knl/knl_variable.h"

#include "commands/explain.h"
#include "executor/instrument.h"
#include "utils/guc.h"
#include "instruments/instr_statement.h"

PG_MODULE_MAGIC;

/* GUC variables */
static int auto_explain_log_min_duration = -1; /* msec or -1 */
static bool auto_explain_log_analyze = false;
static bool auto_explain_log_verbose = false;
static bool auto_explain_log_buffers = false;
static bool auto_explain_log_timing = false;
static int auto_explain_log_format = EXPLAIN_FORMAT_TEXT;
static bool auto_explain_log_nested_statements = false;

static const struct config_enum_entry format_options[] = {{"text", EXPLAIN_FORMAT_TEXT, false},
    {"xml", EXPLAIN_FORMAT_XML, false},
    {"json", EXPLAIN_FORMAT_JSON, false},
    {"yaml", EXPLAIN_FORMAT_YAML, false},
    {NULL, 0, false}};

/* Current nesting depth of ExecutorRun calls */
static int nesting_level = 0;

/* Saved hook values in case of unload */
static ExecutorStart_hook_type prev_ExecutorStart = NULL;
static ExecutorRun_hook_type prev_ExecutorRun = NULL;
static ExecutorFinish_hook_type prev_ExecutorFinish = NULL;
static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;

#define auto_explain_enabled() \
    (auto_explain_log_min_duration >= 0 && (nesting_level == 0 || auto_explain_log_nested_statements))

extern "C" void _PG_init(void);
extern "C" void _PG_fini(void);

static void explain_ExecutorStart(QueryDesc* queryDesc, int eflags);
static void explain_ExecutorRun(QueryDesc* queryDesc, ScanDirection direction, long count);
static void explain_ExecutorFinish(QueryDesc* queryDesc);
static void explain_ExecutorEnd(QueryDesc* queryDesc);

/*
 * Module load callback
 */
void _PG_init(void)
{
    /* Define custom GUC variables. */
    DefineCustomIntVariable("auto_explain.log_min_duration",
        "Sets the minimum execution time above which plans will be logged.",
        "Zero prints all plans. -1 turns this feature off.",
        &auto_explain_log_min_duration,
        -1,
        -1,
        INT_MAX / 1000,
        PGC_SUSET,
        GUC_UNIT_MS,
        NULL,
        NULL,
        NULL);

    DefineCustomBoolVariable("auto_explain.log_analyze",
        "Use EXPLAIN ANALYZE for plan logging.",
        NULL,
        &auto_explain_log_analyze,
        false,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    DefineCustomBoolVariable("auto_explain.log_verbose",
        "Use EXPLAIN VERBOSE for plan logging.",
        NULL,
        &auto_explain_log_verbose,
        false,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    DefineCustomBoolVariable("auto_explain.log_buffers",
        "Log buffers usage.",
        NULL,
        &auto_explain_log_buffers,
        false,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    DefineCustomEnumVariable("auto_explain.log_format",
        "EXPLAIN format to be used for plan logging.",
        NULL,
        &auto_explain_log_format,
        EXPLAIN_FORMAT_TEXT,
        format_options,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    DefineCustomBoolVariable("auto_explain.log_nested_statements",
        "Log nested statements.",
        NULL,
        &auto_explain_log_nested_statements,
        false,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    DefineCustomBoolVariable("auto_explain.log_timing",
        "Collect timing data, not just row counts.",
        NULL,
        &auto_explain_log_timing,
        true,
        PGC_SUSET,
        0,
        NULL,
        NULL,
        NULL);

    EmitWarningsOnPlaceholders("auto_explain");

    /* Install hooks. */
    prev_ExecutorStart = ExecutorStart_hook;
    ExecutorStart_hook = explain_ExecutorStart;
    prev_ExecutorRun = ExecutorRun_hook;
    ExecutorRun_hook = explain_ExecutorRun;
    prev_ExecutorFinish = ExecutorFinish_hook;
    ExecutorFinish_hook = explain_ExecutorFinish;
    prev_ExecutorEnd = ExecutorEnd_hook;
    ExecutorEnd_hook = explain_ExecutorEnd;
}

/*
 * Module unload callback
 */
void _PG_fini(void)
{
    /* Uninstall hooks. */
    ExecutorStart_hook = prev_ExecutorStart;
    ExecutorRun_hook = prev_ExecutorRun;
    ExecutorFinish_hook = prev_ExecutorFinish;
    ExecutorEnd_hook = prev_ExecutorEnd;
}

/*
 * ExecutorStart hook: start up logging if needed
 */
static void explain_ExecutorStart(QueryDesc* queryDesc, int eflags)
{
    if (auto_explain_enabled()) {
        /* Enable per-node instrumentation iff log_analyze is required. */
        if (auto_explain_log_analyze && (eflags & EXEC_FLAG_EXPLAIN_ONLY) == 0) {
            if (auto_explain_log_timing)
                queryDesc->instrument_options |= INSTRUMENT_TIMER;
            else
                queryDesc->instrument_options |= INSTRUMENT_ROWS;

            if (auto_explain_log_buffers)
                queryDesc->instrument_options |= INSTRUMENT_BUFFERS;
        }
    }

    if (prev_ExecutorStart)
        prev_ExecutorStart(queryDesc, eflags);
    else
        standard_ExecutorStart(queryDesc, eflags);

    if (auto_explain_enabled()) {
        /*
         * Set up to track total elapsed time in ExecutorRun.  Make sure the
         * space is allocated in the per-query context so it will go away at
         * ExecutorEnd.
         */
        if (queryDesc->totaltime == NULL) {
            MemoryContext oldcxt;

            oldcxt = MemoryContextSwitchTo(queryDesc->estate->es_query_cxt);
            queryDesc->totaltime = InstrAlloc(1, INSTRUMENT_ALL);
            MemoryContextSwitchTo(oldcxt);
        }
    }
}

/*
 * ExecutorRun hook: all we need do is track nesting depth
 */
static void explain_ExecutorRun(QueryDesc* queryDesc, ScanDirection direction, long count)
{
    nesting_level++;
    PG_TRY();
    {
        if (prev_ExecutorRun)
            prev_ExecutorRun(queryDesc, direction, count);
        else
            standard_ExecutorRun(queryDesc, direction, count);
        nesting_level--;
    }
    PG_CATCH();
    {
        nesting_level--;
        instr_stmt_exec_report_query_plan(queryDesc);
        PG_RE_THROW();
    }
    PG_END_TRY();
}

/*
 * ExecutorFinish hook: all we need do is track nesting depth
 */
static void explain_ExecutorFinish(QueryDesc* queryDesc)
{
    nesting_level++;
    PG_TRY();
    {
        if (prev_ExecutorFinish)
            prev_ExecutorFinish(queryDesc);
        else
            standard_ExecutorFinish(queryDesc);
        nesting_level--;
    }
    PG_CATCH();
    {
        nesting_level--;
        PG_RE_THROW();
    }
    PG_END_TRY();
}

/*
 * ExecutorEnd hook: log results if needed
 */
static void explain_ExecutorEnd(QueryDesc* queryDesc)
{
    if (queryDesc->totaltime && auto_explain_enabled()) {
        double msec;

        /*
         * Make sure stats accumulation is done.  (Note: it's okay if several
         * levels of hook all do this.)
         */
        InstrEndLoop(queryDesc->totaltime);

        /* Log plan if duration is exceeded. */
        msec = queryDesc->totaltime->total * 1000.0;
        if (msec >= auto_explain_log_min_duration) {
            ExplainState es;

            ExplainInitState(&es);
            es.analyze = (queryDesc->instrument_options && auto_explain_log_analyze);
            es.verbose = auto_explain_log_verbose;
            es.buffers = (es.analyze && auto_explain_log_buffers);
            es.format = (ExplainFormat)auto_explain_log_format;

            ExplainBeginOutput(&es);
            ExplainQueryText(&es, queryDesc);
            ExplainPrintPlan(&es, queryDesc);
            ExplainEndOutput(&es);

            /* Remove last line break */
            if (es.str->len > 0 && es.str->data[es.str->len - 1] == '\n')
                es.str->data[--es.str->len] = '\0';

            /* Fix JSON to output an object */
            if (auto_explain_log_format == EXPLAIN_FORMAT_JSON) {
                es.str->data[0] = '{';
                es.str->data[es.str->len - 1] = '}';
            }

            /*
             * Note: we rely on the existing logging of context or
             * debug_query_string to identify just which statement is being
             * reported.  This isn't ideal but trying to do it here would
             * often result in duplication.
             */
            ereport(LOG, (errmsg("duration: %.3f ms  plan:\n%s", msec, es.str->data), errhidestmt(true)));

            pfree(es.str->data);
        }
    }

    if (prev_ExecutorEnd)
        prev_ExecutorEnd(queryDesc);
    else
        standard_ExecutorEnd(queryDesc);
}
